"""数据加载与处理模块.

本模块负责从各个数据源加载和处理数据:
1. 数据库数据加载
2. CSV文件处理
3. 评分数据存储
"""

from pathlib import Path
from typing import Dict, List, Optional, Union

import pandas as pd
import polars as pl
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine

from ..utils.config import Config
from ..utils.logger import logger


class DataLoader:
    """数据加载器类,用于处理各类数据源的数据加载和处理."""
    
    def __init__(self):
        """初始化数据加载器."""
        self.engine: Optional[Engine] = None
    
    def get_engine(self) -> Engine:
        """获取或创建数据库引擎.
        
        Returns:
            Engine: SQLAlchemy数据库引擎实例
        """
        if self.engine is None:
            self.engine = create_engine(Config.get_database_url())
        return self.engine
    
    def load_merchant_data(self, merchant_id: str) -> Dict[str, float]:
        """加载商家数据用于评分.
        
        加载以下维度数据:
        - 交易金额和次数
        - 服务评级
        - 商品质量评分
        - 平台活跃度
        - 合规分数
        
        Args:
            merchant_id: 商家ID
            
        Returns:
            Dict[str, float]: 商家特征数据字典
        """
        try:
            query = """
            SELECT 
                transaction_amount,    -- 交易总额
                transaction_count,     -- 交易次数
                service_rating,        -- 服务评分
                product_quality,       -- 商品质量
                platform_activity,     -- 平台活跃度
                compliance_score       -- 合规分数
            FROM merchant_features
            WHERE merchant_id = :merchant_id
            """
            
            df = pd.read_sql_query(
                query,
                self.get_engine(),
                params={"merchant_id": merchant_id}
            )
            
            if df.empty:
                logger.warning(f"未找到商家数据: {merchant_id}")
                return {}
                
            return df.iloc[0].to_dict()
            
        except Exception as e:
            logger.error(f"加载商家数据时发生错误: {e}")
            raise
    
    def load_user_data(self, user_id: str) -> Dict[str, float]:
        """加载用户数据用于评分.
        
        加载以下维度数据:
        - 消费金额
        - 交易频率
        - 支付可靠性
        - 退货率
        - 投诉率
        - 社交评分
        
        Args:
            user_id: 用户ID
            
        Returns:
            Dict[str, float]: 用户特征数据字典
        """
        try:
            query = """
            SELECT 
                consumption_amount,     -- 消费总额
                transaction_frequency,  -- 交易频率
                payment_reliability,    -- 支付可靠性
                return_rate,           -- 退货率
                complaint_rate,        -- 投诉率
                social_score           -- 社交评分
            FROM user_features
            WHERE user_id = :user_id
            """
            
            df = pd.read_sql_query(
                query,
                self.get_engine(),
                params={"user_id": user_id}
            )
            
            if df.empty:
                logger.warning(f"未找到用户数据: {user_id}")
                return {}
                
            return df.iloc[0].to_dict()
            
        except Exception as e:
            logger.error(f"加载用户数据时发生错误: {e}")
            raise
    
    def load_transaction_data(
        self,
        start_date: str,
        end_date: str,
        merchant_id: Optional[str] = None,
        user_id: Optional[str] = None
    ) -> pd.DataFrame:
        """加载交易数据用于风险评估.
        
        Args:
            start_date: 开始日期
            end_date: 结束日期
            merchant_id: 可选的商家ID过滤
            user_id: 可选的用户ID过滤
            
        Returns:
            pd.DataFrame: 交易数据DataFrame
        """
        try:
            query = """
            SELECT *
            FROM transactions
            WHERE transaction_date BETWEEN :start_date AND :end_date
            """
            
            params = {
                "start_date": start_date,
                "end_date": end_date
            }
            
            # 添加商家ID过滤
            if merchant_id:
                query += " AND merchant_id = :merchant_id"
                params["merchant_id"] = merchant_id
            
            # 添加用户ID过滤
            if user_id:
                query += " AND user_id = :user_id"
                params["user_id"] = user_id
            
            return pd.read_sql_query(query, self.get_engine(), params=params)
            
        except Exception as e:
            logger.error(f"加载交易数据时发生错误: {e}")
            raise
    
    def save_score(
        self,
        entity_id: str,
        entity_type: str,
        score: float,
        score_details: Dict[str, float]
    ) -> bool:
        """保存信用评分和详细得分.
        
        Args:
            entity_id: 实体ID(商家或用户)
            entity_type: 实体类型('merchant'或'user')
            score: 总体信用评分
            score_details: 各维度详细得分
            
        Returns:
            bool: 保存是否成功
        """
        try:
            # 将评分详情转换为JSON格式
            score_json = pd.Series(score_details).to_json()
            
            # 构建UPSERT语句
            query = """
            INSERT INTO credit_scores (
                entity_id,           -- 实体ID
                entity_type,         -- 实体类型
                score,               -- 总分
                score_details,       -- 详细得分
                created_at           -- 创建时间
            ) VALUES (
                :entity_id,
                :entity_type,
                :score,
                :score_details,
                CURRENT_TIMESTAMP
            )
            ON CONFLICT (entity_id, entity_type)
            DO UPDATE SET
                score = :score,
                score_details = :score_details,
                updated_at = CURRENT_TIMESTAMP
            """
            
            with self.get_engine().connect() as conn:
                conn.execute(
                    query,
                    {
                        "entity_id": entity_id,
                        "entity_type": entity_type,
                        "score": score,
                        "score_details": score_json
                    }
                )
                conn.commit()
            
            return True
            
        except Exception as e:
            logger.error(f"保存评分数据时发生错误: {e}")
            return False
    
    @staticmethod
    def load_csv(
        file_path: Union[str, Path],
        use_polars: bool = False
    ) -> Union[pd.DataFrame, pl.DataFrame]:
        """从CSV文件加载数据.
        
        Args:
            file_path: CSV文件路径
            use_polars: 是否使用Polars替代Pandas
            
        Returns:
            Union[pd.DataFrame, pl.DataFrame]: 加载的数据框
        """
        try:
            file_path = Path(file_path)
            if not file_path.exists():
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            if use_polars:
                return pl.read_csv(file_path)
            return pd.read_csv(file_path)
            
        except Exception as e:
            logger.error(f"加载CSV文件时发生错误: {e}")
            raise